from dataclasses import dataclass
from datetime import datetime
from pathlib import Path

import tyro
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.vec_env import DummyVecEnv


@dataclass
class Args:
    v2g: bool = True
    timesteps: int = 500_0000
    seed: int = 42
    envs: int = 8


args = tyro.cli(Args)
venv = make_vec_env(
    "envs:Home-v4",
    n_envs=args.envs,
    vec_env_cls=DummyVecEnv,
    env_kwargs={
        "auto_normalize": True,
        "ev": {"v2g": args.v2g, "enable": True},
        "carbon_emission": {"enable": True},
    },
)

time_str = datetime.now().strftime("%Y%m%d_%H%M%S")
save_path = Path("output") / time_str

# save callback
checkpoint_callback = CheckpointCallback(
    save_freq=20000,
    save_path=save_path / "checkpoints",
    name_prefix="model",
)

model = PPO("MlpPolicy", venv, tensorboard_log=save_path)
model.learn(
    total_timesteps=args.timesteps,
    progress_bar=True,
    log_interval=1,
    callback=checkpoint_callback,
)
